#!/usr/bin/env python3
import asyncio
import logging
# import sys
from typing import Dict, Any
from mcp.server import Server
from mcp.types import Resource, Tool, TextContent
from mcp.server.stdio import stdio_server

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("stk_mcp_server")

# 创建MCP应用实例
app = Server("stk_mcp_server")

from stk_interface import load_config, create_scenario_inter, exec_task_inter
# 初始化STK应用
class StkApplication:
    def __init__(self):
        self.obj_dic:Dict[str, Any] = {}
        # self.scenario, self.analysis_list = load_config()
        pass

    def create_scenario(self, config_scenario: Dict[str, Any]) -> Dict[str, Any]:
        self.obj_dic = create_scenario_inter(config_scenario)
        return {"result":str(self.obj_dic)}

    def execute_task(self, params: Dict[str, Any]) -> Dict[str, Any]:
        # 定义一个方法execute_task，接收一个参数params，params是一个字典，键为字符串，值为任意类型
        # 返回一个字典，键为字符串，值为任意类型
        # 将params转换为字符串，并将其作为值，键为"result"，返回这个字典
        # res = exec_task_inter(self.analysis_list, self.obj_dic)
        res = exec_task_inter(params["analysis_list"], self.obj_dic)
        return {"result":str(res)}
    
    @staticmethod
    def Start():
        return StkApplication()
    
stk_app = StkApplication.Start()

@app.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="create_scenario",
            description="Create a new STK scenario",
            inputSchema={
                "type": "object",
                "properties": {
                    "config_scenario": {
                        "type": "object",
                        "required": ["objects"],
                        "properties": {
                            "scenario_name": {
                                "type":"string",
                                "description": "Name of the scenario"
                            },
                            "objects": {
                                "type": "array",
                                "minItems": 1,
                                "items": {
                                    "type": "object",
                                    "required": ["type", "name", "basic"],
                                    "properties": {
                                    "type": {
                                        "type": "string",
                                        "description": "Type of the object"
                                    },
                                    "name": {
                                        "type": "string",
                                        "description": "Name of the object"
                                    },
                                    "basic": {
                                        "type": "object"
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        ),
        Tool(
            name="execute_task",
            description="Execute STK task",
            inputSchema={
                "type": "object",
                "required": ["analysis_list", "scenario_objs"],
                "properties": {
                    "analysis_list": {
                        "type": "array",
                        "description": "List of analysis to execute",
                        "minItems": 1,
                        "items": {
                            "type": "object",
                            "required": ["analysis_type", "parameters"],
                            "properties": {
                                "analysis_type": {
                                    "type": "string",
                                    "description": "analysis's type to execute"
                                },
                                "parameters": {
                                    "type": "object",
                                    "description": "analysis's parameters"
                                },
                                "oneOf": [
                                    {
                                        "properties": {
                                        "analysis_type": {
                                            "const": "Access"
                                        },
                                        "parameters": {
                                            "type": "object",
                                            "required": ["observed", "observer"],
                                            "description": "需要计算Access的2个对象",
                                            "properties": {
                                            "observed": {
                                                "type": "string"
                                            },
                                            "observer": {
                                                "type": "string"
                                            }
                                            }
                                        }
                                        }
                                    },
                                    {
                                        "properties": {
                                        "analysis_type": {
                                            "not": {
                                            "const": "Access"
                                            }
                                        }
                                        }
                                    }
                                ]
                            }
                        }
                    },
                    "scenario_objs": {
                        "type": "object",
                        "description": "The Scenarios objects Created by function 'create_scenario'."
                    }
                }
            }
        )
    ]

@app.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    logger.info(f"Calling tool: {name} with arguments: {arguments}")
    
    if name == "create_scenario":
        config_scenario = arguments.get("config_scenario")
        if not config_scenario:
            return [TextContent(type="text", text="Error: config_scenario is required")]
            
        try:
            result = stk_app.create_scenario(config_scenario)
            return [TextContent(
                type="text",
                text=f"Scenario '{result['result']}' created successfully"
            )]
        except Exception as e:
            return [TextContent(
                type="text",
                text=f"Error creating scenario: {str(e)}"
            )]
            
    elif name == "execute_task":
        analysis_list, scenario_objs = arguments.get("analysis_list"), arguments.get("scenario_objs")
        if not analysis_list or not scenario_objs:
            return [TextContent(type="text", text="Error: analysis_list and scenario_objs is required")]
            
        try:
            # 这里实现具体的任务执行逻辑
            result = stk_app.execute_task(arguments)
            return [TextContent(type="text", text=result['result'])]
        except Exception as e:
            return [TextContent(
                type="text", 
                text=f"Error executing task: {str(e)}"
            )]
            
    else:
        return [TextContent(
            type="text",
            text=f"Unknown tool: {name}"
        )]

async def main():
    logger.info("Starting STK MCP server...")
    async with stdio_server() as (read_stream, write_stream):
        try:
            await app.run(
                read_stream,
                write_stream,
                app.create_initialization_options()
            )
        except Exception as e:
            logger.error(f"Server error: {str(e)}", exc_info=True)
            raise

if __name__ == "__main__":
    asyncio.run(main())
